import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import matplotlib as mpl
import matplotlib.font_manager as fm

# 设置支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

# 设置参数
a = 1.0       # 扩散系数
A = 1.0       # 振幅
omega = 4.0   # 角频率
l = np.pi     # 空间长度
N_terms = 50  # 级数项数

# 创建空间和时间网格
x = np.linspace(0, l, 200)
t_max = 5.0
t_points = np.linspace(0, t_max, 100)

# 预计算特征值
n_values = np.arange(0, N_terms)
k_n = (n_values + 0.5) * np.pi / l
lambda_n = (k_n * a)**2

# 计算解析解
def u(x, t):
    result = np.zeros_like(x)
    for n in range(N_terms):
        # 计算系数
        denom = lambda_n[n]**2 + omega**2
        coeff = (2 * A / (np.pi * (n + 0.5))) / denom
        
        # 计算时间相关部分
        trig_part = lambda_n[n] * np.sin(omega*t) - omega * np.cos(omega*t)
        exp_part = omega * np.exp(-lambda_n[n] * t)
        time_part = trig_part + exp_part
        
        # 空间部分
        space_part = np.sin((n + 0.5) * np.pi * x / l)
        
        # 累加到结果
        result += coeff * time_part * space_part
    return result

# 创建图形
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10), gridspec_kw={'height_ratios': [1, 3]})
fig.subplots_adjust(hspace=0.4)

# 设置颜色映射 - 蓝白红渐变
colors = ["blue", "cyan", "white", "yellow", "red"]
cmap = mpl.colors.ListedColormap(colors)
norm = mpl.colors.Normalize(vmin=-0.5, vmax=0.5)

# 上子图：温度曲线
ax1.set_xlim(0, l)
ax1.set_ylim(-0.6, 0.6)
ax1.set_xlabel('位置 x')
ax1.set_ylabel('温度 u(x,t)')
ax1.set_title('细长杆温度分布')
ax1.grid(True, alpha=0.3)

line, = ax1.plot([], [], lw=2, color='black')
time_text = ax1.text(0.02, 0.95, '', transform=ax1.transAxes, fontsize=12)

# 下子图：温度热图
ax2.set_xlim(0, l)
ax2.set_ylim(0, 0.2)  # 杆的厚度
ax2.set_xlabel('位置 x')
ax2.set_ylabel('')
ax2.set_title('温度分布热图')
ax2.set_yticks([])  # 隐藏y轴刻度

# 创建热图网格
X, Y = np.meshgrid(x, np.linspace(0, 0.2, 10))
# 初始温度值全为零
Z = np.zeros_like(X)

# 创建热图 - 使用正确的数组形状
heatmap = ax2.pcolormesh(X, Y, Z, cmap=cmap, norm=norm, shading='auto')
cbar = fig.colorbar(heatmap, ax=ax2)
cbar.set_label('温度值')

# 初始化函数
def init():
    line.set_data([], [])
    time_text.set_text('')
    # 设置初始温度值
    Z_init = np.zeros((10, 200))
    heatmap.set_array(Z_init.ravel())
    return line, heatmap, time_text

# 动画更新函数
def update(frame):
    t = t_points[frame]
    u_vals = u(x, t)
    
    # 更新曲线
    line.set_data(x, u_vals)
    time_text.set_text(f'时间 t = {t:.2f}')
    
    # 更新热图 - 创建与网格匹配的2D温度数组
    Z = np.tile(u_vals, (10, 1))  # 创建2D数组表示杆的温度
    
    # 确保数组大小匹配网格
    if Z.shape != (10, 200):
        Z = np.zeros((10, 200))  # 如果形状不匹配，使用零数组
        
    heatmap.set_array(Z.ravel())
    
    return line, heatmap, time_text

# 创建动画 - 使用frames参数确保在范围内
ani = FuncAnimation(fig, update, frames=len(t_points), 
                    init_func=init, blit=False, interval=50)

# 保存为GIF - 使用更可靠的保存方式
try:
    writer = PillowWriter(fps=15)
    ani.save('热传导8-2例题3.gif', writer=writer)
    print("动画已成功保存为 '热传导8-2例题3.gif'")
except Exception as e:
    print(f"保存动画时出错: {e}")

# 显示动画
plt.tight_layout()
plt.show()